package server;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

/**
 * The IPTablesComm class is a wrapper around IPTables that allows addition
 * and removal of rules.
 * It uses shell command execution to invoke IPTables commands.
 */
public class IPTablesComm {

    //@{
    /** Constants used for running IPTables */
	static String shell = "/bin/sh";
	static String shellOpt = "-c";
	static final String iptableTable = "mangle";
	static final String iptablesFile = "/sbin/iptables"; 
    //@}

    /**
     * Flushes the IPTables rule table by executing iptables -F on our table.
     *
     * @throw FirewallException  on any error
     */
	void flush() throws FirewallException {
		// flush default allow/accept is allow?
		run(iptablesFile + " -t " + iptableTable + " -F");
	}

    /**
     * Adds an IPTables rule with the given parameters.
     * The arguments bear slightly different meanings depending on the type
     * of the rule:
     * - For IP rules, logging means that the Firewall will log packets that
     *   matched the rule. The regexp argument must be NULL for IP rules.
     * - For context rules, logging means that the Firewall will log all
     *   sessions that the rule positively identified or rejected. The regexp
     *   argument must be non-NULL for context rules.
     *
     * Note that logging of IP rules is supported by adding two rules: one for
     * logging and one for actually filtering the data. Context rules,
     * however, support their own logging.
     *
     * @param   dir     The direction of the rule (INPUT or OUTPUT)
     * @param   rule    The rule itself
     * @param   log     Should this rule log events. See method documentation
     *                  for details.
     * @param   regexp  The regexp to use for the rule (applicable for
     *                  context rules only, specify NULL for IP rules)
     *
     * @throw FirewallException  on any error
     */
	void add(RulesTableType dir, Rule rule, boolean log, String regexp)
        throws FirewallException {

        // Use 'dir' to prepare the direction string we'll pass to IPTables.
		String Direction;
		if (dir == RulesTableType.IN)
			Direction = "INPUT";
		else
			Direction = "OUTPUT";

		if (rule.type == Rule.typeOpt.CONTEXT)
		{
            //
            // Context rule
            //
			run(
                    iptablesFile + 
                    " -t " + iptableTable + 
                    " -I " + Direction + " " + 
                    rule.index +
                    " -" + ((dir == RulesTableType.IN) ? "i" : "o") + " ! lo "
                    + " -m ipcontext --ipcontext_port " + rule.port + 
                    " --ipcontext_protocol " + rule.protocol + 
                    ((dir == RulesTableType.IN) ? " --ipcontext_dir 1 " : "") + 
                    (log ? " --ipcontext_log " : "") + 
                    " --ipcontext_regexp \"" + regexp + 
                    "\" -j " + ((rule.permission.toString().compareTo("ALLOW") == 0) ? 
                                "ACCEPT" : "DROP"));
		} else {
            //
            // IP rule
            //
			String ruler = "";

            // Deal with IP address range if we were asked to.
			if (rule.iprange != null)
			{
				boolean range = rule.iprange.split("-").length == 2;
				String prefix = 
                    (dir == RulesTableType.IN) ? 
                        (range ? "-src-range" : "s") : 
                        (range ? "-dst-range" : "d");
				ruler += " -" + prefix + " " + rule.iprange + " ";
			}

            // Add port if it's present
			if (rule.port != null && rule.port != 0)
			{
				String prefix = (dir == RulesTableType.IN) ? "-sport" : "-dport";
				ruler += " -" + prefix + " " + rule.port + " ";				
			}

            // Add the actual filtering rule.
			run(
                    iptablesFile + 
                    " -t " + iptableTable + 
                    " -I " + Direction + " " + 
                    rule.index + 
                    " -" + ((dir == RulesTableType.IN) ? "i" : "o") + " ! lo "
                    + " -p tcp " + ruler + 
                    " -j " + ((rule.permission.toString().compareTo("ALLOW") == 0) ? 
                        "ACCEPT" : 
                        "DROP")); 

            // If we should log, add an extra rule just for logging.
			if (log)
				run(
                        iptablesFile + 
                        " -t " + iptableTable + 
                        " -I " + Direction + " " + 
                        rule.index + 
                        " -" + ((dir == RulesTableType.IN) ? "i" : "o") + " ! lo "
                        + " -p tcp " + ruler + 
                        " -j LOG --log-prefix \"OSLFW:\" ");
		}
	}

    /**
     * Removes an IPTables rule from the given table and index.
     *
     * @param   dir     The direction to remove from (INPUT or OUTPUT)
     * @param   index   Index of the rule to remove
     *
     * @throw FirewallException on any error (e.g., invalid index)
     */
	void remove(RulesTableType dir, int index) throws FirewallException {

        // Use 'dir' to prepare the direction string we'll pass to IPTables.
		String Direction;
		if (dir == RulesTableType.IN)
			Direction = "INPUT";
		else
			Direction = "OUTPUT";

        // Remove the rule
		run(
                iptablesFile + 
                " -t " + iptableTable + 
                " -D " + Direction + " " + 
                index);
    }

    /**
     * Executes the given command using shell and shellOpt (e.g., "/bin/sh -c ls").
     * The command's standard output and standard error are printed.
     *
     * @param   command The command to run
     *
     * @throw   FirewallException on any error, or if the command returned
     *                            a return value different than 0.
     */
	private void run(String command) throws FirewallException
	{
		try {
            //
            // Run the command
            //
			String errorMsg = "";
			String[] cmd = {shell, shellOpt, command};
			System.out.println("Running: " + command);
			Process proc = Runtime.getRuntime().exec(cmd);

            //
            // Print out the standard output and standard error
            //
			BufferedReader input = 
                new BufferedReader(new InputStreamReader(proc.getInputStream()));
			BufferedReader error = 
                new BufferedReader(new InputStreamReader(proc.getErrorStream()));
			String line;
			while ((line = input.readLine()) != null)
			{
				System.out.println(line);
			}
			while ((line = error.readLine()) != null)
			{
				errorMsg += line + "\n";
				System.out.println(line);
			}
			input.close();

            // Examine the return value. If it's not 0 then the command
            // failed and we throw an exception.
			int retval;
			try {
				retval = proc.waitFor();
				System.out.println("Return value was: " + retval);
				if (retval != 0)
					throw new FirewallException(errorMsg);
			} catch (InterruptedException e) {
			}
		} catch (IOException e) {
			throw new FirewallException();
		}
	}
}
